import cv2
import numpy as np

# 初始化摄像头
video_capture = cv2.VideoCapture(0)

# 初始化人数统计变量
total_people = 0

# 加载预训练模型（需下载模型文件和配置文件）
model_weights = "yolov3.weights"
model_config = "yolov3.cfg"
net = cv2.dnn.readNetFromDarknet(model_config, model_weights)
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)

# 定义类别标签（COCO数据集中'person'类别ID为0）
classes = []
with open("coco.names", "r") as f:
    classes = [line.strip() for line in f.readlines()]
person_class_id = 0  # 假设检测人体，通过框的上半部分近似人头

# 或使用专门的人头检测模型（需自定义类别）

while True:
    ret, frame = video_capture.read()
    if not ret:
        print("无法获取视频帧")
        break

    # 构造输入Blob
    blob = cv2.dnn.blobFromImage(
        frame, 1/255.0, (416, 416), swapRB=True, crop=False
    )
    net.setInput(blob)

    # 获取检测结果
    output_layers = net.getUnconnectedOutLayersNames()
    detections = net.forward(output_layers)

    # 解析检测结果
    h, w = frame.shape[:2]
    boxes = []
    confidences = []

    for out in detections:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]

            # 只保留置信度高的检测结果（阈值可调）
            if confidence > 0.5 and class_id == person_class_id:
                # 获取边界框坐标
                box = detection[0:4] * np.array([w, h, w, h])
                (center_x, center_y, width, height) = box.astype("int")
                x = int(center_x - (width / 2))
                y = int(center_y - (height / 2))

                # 近似人头位置（取人体框上半部分）
                head_height = int(height * 0.4)  # 调整比例
                head_box = (x, y, x + width, y + head_height)
                boxes.append(head_box)
                confidences.append(float(confidence))

    # 非极大值抑制（NMS）去重
    indices = cv2.dnn.NMSBoxes(
        boxes, confidences, 0.5, 0.4
    )

    # 统计当前人数
    current_people = len(indices)
    if current_people > total_people:
        total_people = current_people

    # 绘制人头框
    for i in indices.flatten():
        (x1, y1, x2, y2) = boxes[i]
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

    # 显示统计信息
    cv2.putText(frame, f"Current: {current_people}", (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
    cv2.putText(frame, f"Peak: {total_people}", (10, 60),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)

    cv2.imshow('Head Detection', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

video_capture.release()
cv2.destroyAllWindows()